{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nCompile PyTorch Models\n======================\n**Author**: `Alex Wong <https://github.com/alexwong/>`_\n\nThis article is an introductory tutorial to deploy PyTorch models with Relay.\n\nFor us to begin with, PyTorch should be installed.\nTorchVision is also required since we will be using it as our model zoo.\n\nA quick solution is to install via pip\n\n.. code-block:: bash\n\n    pip install torch==1.7.0\n    pip install torchvision==0.8.1\n\nor please refer to official site\nhttps://pytorch.org/get-started/locally/\n\nPyTorch versions should be backwards compatible but should be used\nwith the proper TorchVision version.\n\nCurrently, TVM supports PyTorch 1.7 and 1.4. Other versions may\nbe unstable.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tvm\nfrom tvm import relay\n\nimport numpy as np\n\nfrom tvm.contrib.download import download_testdata\n\n# PyTorch imports\nimport torch\nimport torchvision"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load a pretrained PyTorch model\n-------------------------------\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model_name = \"resnet18\"\nmodel = getattr(torchvision.models, model_name)(pretrained=True)\nmodel = model.eval()\n\n# We grab the TorchScripted model via tracing\ninput_shape = [1, 3, 224, 224]\ninput_data = torch.randn(input_shape)\nscripted_model = torch.jit.trace(model, input_data).eval()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Load a test image\n-----------------\nClassic cat example!\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n\nimg_url = \"https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true\"\nimg_path = download_testdata(img_url, \"cat.png\", module=\"data\")\nimg = Image.open(img_path).resize((224, 224))\n\n# Preprocess the image and convert to tensor\nfrom torchvision import transforms\n\nmy_preprocess = transforms.Compose(\n    [\n        transforms.Resize(256),\n        transforms.CenterCrop(224),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ]\n)\nimg = my_preprocess(img)\nimg = np.expand_dims(img, 0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Import the graph to Relay\n-------------------------\nConvert PyTorch graph to Relay graph. The input name can be arbitrary.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "input_name = \"input0\"\nshape_list = [(input_name, img.shape)]\nmod, params = relay.frontend.from_pytorch(scripted_model, shape_list)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Relay Build\n-----------\nCompile the graph to llvm target with given input specification.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "target = tvm.target.Target(\"llvm\", host=\"llvm\")\ndev = tvm.cpu(0)\nwith tvm.transform.PassContext(opt_level=3):\n    lib = relay.build(mod, target=target, params=params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Execute the portable graph on TVM\n---------------------------------\nNow we can try deploying the compiled model on target.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from tvm.contrib import graph_executor\n\ndtype = \"float32\"\nm = graph_executor.GraphModule(lib[\"default\"](dev))\n# Set inputs\nm.set_input(input_name, tvm.nd.array(img.astype(dtype)))\n# Execute\nm.run()\n# Get outputs\ntvm_output = m.get_output(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Look up synset name\n-------------------\nLook up prediction top 1 index in 1000 class synset.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "synset_url = \"\".join(\n    [\n        \"https://raw.githubusercontent.com/Cadene/\",\n        \"pretrained-models.pytorch/master/data/\",\n        \"imagenet_synsets.txt\",\n    ]\n)\nsynset_name = \"imagenet_synsets.txt\"\nsynset_path = download_testdata(synset_url, synset_name, module=\"data\")\nwith open(synset_path) as f:\n    synsets = f.readlines()\n\nsynsets = [x.strip() for x in synsets]\nsplits = [line.split(\" \") for line in synsets]\nkey_to_classname = {spl[0]: \" \".join(spl[1:]) for spl in splits}\n\nclass_url = \"\".join(\n    [\n        \"https://raw.githubusercontent.com/Cadene/\",\n        \"pretrained-models.pytorch/master/data/\",\n        \"imagenet_classes.txt\",\n    ]\n)\nclass_name = \"imagenet_classes.txt\"\nclass_path = download_testdata(class_url, class_name, module=\"data\")\nwith open(class_path) as f:\n    class_id_to_key = f.readlines()\n\nclass_id_to_key = [x.strip() for x in class_id_to_key]\n\n# Get top-1 result for TVM\ntop1_tvm = np.argmax(tvm_output.numpy()[0])\ntvm_class_key = class_id_to_key[top1_tvm]\n\n# Convert input to PyTorch variable and get PyTorch result for comparison\nwith torch.no_grad():\n    torch_img = torch.from_numpy(img)\n    output = model(torch_img)\n\n    # Get top-1 result for PyTorch\n    top1_torch = np.argmax(output.numpy())\n    torch_class_key = class_id_to_key[top1_torch]\n\nprint(\"Relay top-1 id: {}, class name: {}\".format(top1_tvm, key_to_classname[tvm_class_key]))\nprint(\"Torch top-1 id: {}, class name: {}\".format(top1_torch, key_to_classname[torch_class_key]))"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}